#!/usr/bin/env python3

import argparse
import numpy as np
import pandas as pd
import pendulum
import netCDF4
import pymongo
import os
import re
from scipy.interpolate import interp1d, interp2d
import dask.bag
from pprint import pprint

def parse_time(string):
	match = re.match(r'(\d{4}\d{2}\d{2}\d{2})(\d{2})?', string)
	if match.group(2):
		return pendulum.from_format(string, 'YYYYMMDDHHmm')
	else:
		return pendulum.from_format(string, 'YYYYMMDDHH')

parser = argparse.ArgumentParser(description='Calculate OMB.')
parser.add_argument('--host', help='Observation database host', required=True)
parser.add_argument('--port', help='Observation database port', required=True, type=int)
parser.add_argument('--synop', action='store_true')
parser.add_argument('--raob', action='store_true')
parser.add_argument('--amdar', action='store_true')
parser.add_argument('--profiler', action='store_true')
parser.add_argument('--bkg-root', required=True)
parser.add_argument('--start-time', required=True, type=parse_time)
parser.add_argument('--end-time', required=True, type=parse_time)
args = parser.parse_args()

if not os.path.isdir(args.bkg_root):
	print(f'[Error]: {args.bkg_root} does not exist!')
	exit(0)

missing_value = -1e10

def sub(a, b):
	if a != missing_value:
		return a - b
	else:
		return missing_value

if args.synop:
	# Get SYNOP stations.
	client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
	db = client.metobs
	synop_stations = {}
	for station in db.synop_stations.find():
		synop_stations[station['sid']] = {
			'lon': station['location']['coordinates'][0],
			'lat': station['location']['coordinates'][1],
			'z': station['z']
		}
	db.synop_era5_omb.create_index([('sid', 'text'), ('time', 1)], unique=True)
	client.close()

	def calc_synop_omb(time):
		print(f'==> synop {time} ...')
		bkg = netCDF4.Dataset(f'{args.bkg_root}/era5_{time.format("YYYYMMDDHH")}.nc')
		bkg_lon = bkg.variables['longitude']
		bkg_lat = bkg.variables['latitude']
		bkg_u10_interp = interp2d(bkg_lon, bkg_lat, bkg.variables['u10'][0,:,:], kind='linear', copy=False)
		bkg_v10_interp = interp2d(bkg_lon, bkg_lat, bkg.variables['v10'][0,:,:], kind='linear', copy=False)
		bkg_t2m_interp = interp2d(bkg_lon, bkg_lat, bkg.variables['t2m'][0,:,:], kind='linear', copy=False)
		bkg_d2m_interp = interp2d(bkg_lon, bkg_lat, bkg.variables['d2m'][0,:,:], kind='linear', copy=False)
		bkg_sp_interp  = interp2d(bkg_lon, bkg_lat, bkg.variables['sp' ][0,:,:], kind='linear', copy=False)
		bkg.close()
	
		client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
		db = client.metobs
		synop_records = db.synop_records.find({ 'time': time })
	
		synop_omb = []
		for record in synop_records:
			sid = record['sid']
			if not sid in synop_stations:
				#print(f'There is no synop station {sid} in database!')
				continue
			lon = synop_stations[sid]['lon']
			lat = synop_stations[sid]['lat']
			bkg_u10 = bkg_u10_interp(lon, lat) # m s-1
			bkg_v10 = bkg_v10_interp(lon, lat) # m s-1
			bkg_t2m = bkg_t2m_interp(lon, lat) # K
			bkg_d2m = bkg_d2m_interp(lon, lat) # K
			bkg_sp  = bkg_sp_interp (lon, lat) # Pa
			bkg_sh  = specific_humidity_from_dewpoint(bkg_sp, bkg_d2m) # 1
			omb_ua = sub(record['ua'], bkg_u10)
			omb_va = sub(record['va'], bkg_v10)
			omb_ta = sub(record['ta'], bkg_t2m - 273.15)
			omb_p  = sub(record['p' ], bkg_sp / 100.0)
			omb_sh = sub(record['sh'], bkg_sh)
			synop_omb.append({
				'sid': sid,
				'time': time,
				'omb_ua': float(omb_ua),
				'omb_va': float(omb_va),
				'omb_ta': float(omb_ta),
				'omb_p' : float(omb_p ),
				'omb_sh': float(omb_sh)
			})
	
		if len(synop_omb) > 0:
			synop_era5_omb = db.synop_era5_omb
			synop_era5_omb.insert_many(synop_omb)
	
		client.close()
	
	workers = dask.bag.from_sequence(pendulum.period(args.start_time, args.end_time).range('hours', 6), npartitions=20)
	workers.map(calc_synop_omb).compute()

if args.raob:
	# Get RAOB stations.
	client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
	db = client.metobs
	raob_stations = {}
	for station in db.raob_stations.find():
		raob_stations[station['sid']] = {
			'lon': station['location']['coordinates'][0],
			'lat': station['location']['coordinates'][1],
			'z': station['z']
		}
	db.raob_era5_omb.create_index([('sid', 'text'), ('level_type', 1), ('time', 1), ('p', 1)], unique=True)
	client.close()

	def calc_raob_omb(time):
		print(f'==> raob {time} ...')
		bkg = netCDF4.Dataset(f'{args.bkg_root}/era5_{time.format("YYYYMMDDHH")}.nc')
		bkg_lon = bkg.variables['longitude']
		bkg_lat = bkg.variables['latitude']
		bkg_p   = bkg.variables['level'][:]
		bkg_u_hinterp = [interp2d(bkg_lon, bkg_lat, bkg.variables['u'][0,k,:,:], kind='linear', copy=False) for k in range(len(bkg_p))]
		bkg_v_hinterp = [interp2d(bkg_lon, bkg_lat, bkg.variables['v'][0,k,:,:], kind='linear', copy=False) for k in range(len(bkg_p))]
		bkg_t_hinterp = [interp2d(bkg_lon, bkg_lat, bkg.variables['t'][0,k,:,:], kind='linear', copy=False) for k in range(len(bkg_p))]
		bkg_q_hinterp = [interp2d(bkg_lon, bkg_lat, bkg.variables['q'][0,k,:,:], kind='linear', copy=False) for k in range(len(bkg_p))]
		bkg.close()
	
		client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
		db = client.metobs
		raob_records = db.raob_records.find({ 'time': { '$gte': time.subtract(hours=1), '$lte': time.add(hours=1) } })

		raob_omb = []
		for record in raob_records:
			sid = record['sid']
			if not sid in raob_stations:
				#print(f'There is no raob station {sid} in database!')
				continue
			lon = raob_stations[sid]['lon']
			lat = raob_stations[sid]['lat']
			# Mandatory levels
			if 'man_p' in record:
				for ko in range(len(record['man_p'])):
					obs_p = record['man_p'][ko]
					kb = np.abs(bkg_p - obs_p).argmin()
					kb = kb if obs_p - bkg_p[kb] >= 0 else kb - 1
					if kb == bkg_p.size - 1: continue
					bkg_u = interp1d(bkg_p[kb:kb+2], [bkg_u_hinterp[kb](lon, lat)[0], bkg_u_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_v = interp1d(bkg_p[kb:kb+2], [bkg_v_hinterp[kb](lon, lat)[0], bkg_v_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_t = interp1d(np.log(bkg_p[kb:kb+2]), [bkg_t_hinterp[kb](lon, lat)[0], bkg_t_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(np.log(obs_p)) # K
					bkg_q = interp1d(bkg_p[kb:kb+2], [bkg_q_hinterp[kb](lon, lat)[0], bkg_q_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # 1
					omb_ua = sub(record['man_ua'][ko], bkg_u)
					omb_va = sub(record['man_va'][ko], bkg_v)
					omb_ta = sub(record['man_ta'][ko], bkg_t - 273.15)
					omb_sh = sub(record['man_sh'][ko], bkg_q)
					raob_omb.append({
						'sid': sid,
						'time': record['time'],
						'p': obs_p,
						'level_type': 'man',
						'level_idx': ko,
						'omb_ua': float(omb_ua),
						'omb_va': float(omb_va),
						'omb_ta': float(omb_ta),
						'omb_sh': float(omb_sh)
					})
			# Significant temperature levels
			if 'sigt_p' in record:
				for ko in range(len(record['sigt_p'])):
					obs_p = record['sigt_p'][ko]
					kb = np.abs(bkg_p - obs_p).argmin()
					kb = kb if obs_p - bkg_p[kb] >= 0 else kb - 1
					if kb == bkg_p.size - 1: continue
					bkg_u = interp1d(bkg_p[kb:kb+2], [bkg_u_hinterp[kb](lon, lat)[0], bkg_u_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_v = interp1d(bkg_p[kb:kb+2], [bkg_v_hinterp[kb](lon, lat)[0], bkg_v_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_t = interp1d(np.log(bkg_p[kb:kb+2]), [bkg_t_hinterp[kb](lon, lat)[0], bkg_t_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(np.log(obs_p)) # K
					bkg_q = interp1d(bkg_p[kb:kb+2], [bkg_q_hinterp[kb](lon, lat)[0], bkg_q_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # 1
					omb_ua = sub(record['sigt_ua'][ko], bkg_u)
					omb_va = sub(record['sigt_va'][ko], bkg_v)
					omb_ta = sub(record['sigt_ta'][ko], bkg_t - 273.15)
					omb_sh = sub(record['sigt_sh'][ko], bkg_q)
					raob_omb.append({
						'sid': sid,
						'time': record['time'],
						'p': obs_p,
						'level_type': 'sigt',
						'level_idx': ko,
						'omb_ua': float(omb_ua),
						'omb_va': float(omb_va),
						'omb_ta': float(omb_ta),
						'omb_sh': float(omb_sh)
					})
			# Significant wind levels
			if 'sigw_p' in record:
				for ko in range(len(record['sigw_p'])):
					obs_p = record['sigw_p'][ko]
					kb = np.abs(bkg_p - obs_p).argmin()
					kb = kb if obs_p - bkg_p[kb] >= 0 else kb - 1
					if kb == bkg_p.size - 1: continue
					bkg_u = interp1d(bkg_p[kb:kb+2], [bkg_u_hinterp[kb](lon, lat)[0], bkg_u_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_v = interp1d(bkg_p[kb:kb+2], [bkg_v_hinterp[kb](lon, lat)[0], bkg_v_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_t = interp1d(np.log(bkg_p[kb:kb+2]), [bkg_t_hinterp[kb](lon, lat)[0], bkg_t_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(np.log(obs_p)) # K
					bkg_q = interp1d(bkg_p[kb:kb+2], [bkg_q_hinterp[kb](lon, lat)[0], bkg_q_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # 1
					omb_ua = sub(record['sigw_ua'][ko], bkg_u)
					omb_va = sub(record['sigw_va'][ko], bkg_v)
					omb_ta = sub(record['sigw_ta'][ko], bkg_t - 273.15)
					omb_sh = sub(record['sigw_sh'][ko], bkg_q)
					raob_omb.append({
						'sid': sid,
						'time': record['time'],
						'p': obs_p,
						'level_type': 'sigw',
						'level_idx': ko,
						'omb_ua': float(omb_ua),
						'omb_va': float(omb_va),
						'omb_ta': float(omb_ta),
						'omb_sh': float(omb_sh)
					})
			# Tropopause levels
			if 'trop_p' in record:
				for ko in range(len(record['trop_p'])):
					obs_p = record['trop_p'][ko]
					kb = np.abs(bkg_p - obs_p).argmin()
					kb = kb if obs_p - bkg_p[kb] >= 0 else kb - 1
					if kb == bkg_p.size - 1: continue
					bkg_u = interp1d(bkg_p[kb:kb+2], [bkg_u_hinterp[kb](lon, lat)[0], bkg_u_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_v = interp1d(bkg_p[kb:kb+2], [bkg_v_hinterp[kb](lon, lat)[0], bkg_v_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # m s-1
					bkg_t = interp1d(np.log(bkg_p[kb:kb+2]), [bkg_t_hinterp[kb](lon, lat)[0], bkg_t_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(np.log(obs_p)) # K
					bkg_q = interp1d(bkg_p[kb:kb+2], [bkg_q_hinterp[kb](lon, lat)[0], bkg_q_hinterp[kb+1](lon, lat)[0]], kind='linear', copy=False)(obs_p) # 1
					omb_ua = sub(record['trop_ua'][ko], bkg_u)
					omb_va = sub(record['trop_va'][ko], bkg_v)
					omb_ta = sub(record['trop_ta'][ko], bkg_t - 273.15)
					omb_sh = sub(record['trop_sh'][ko], bkg_q)
					raob_omb.append({
						'sid': sid,
						'time': record['time'],
						'p': obs_p,
						'level_type': 'trop',
						'level_idx': ko,
						'omb_ua': float(omb_ua),
						'omb_va': float(omb_va),
						'omb_ta': float(omb_ta),
						'omb_sh': float(omb_sh)
					})
	
		if len(raob_omb) > 0:
			try:
				db.raob_era5_omb.insert_many(raob_omb)
			except pymongo.errors.BulkWriteError as e:
				if e.details['writeErrors'][0]['code'] == 11000: # Duplicate key, just skip it.
					pprint(e.details['writeErrors'][0])
					pass
				else:
					print(f'[Error]: Failed to write database! See error: {e.details["writeErrors"][0]["errmsg"]}')
					exit(1)
	
		client.close()

	workers = dask.bag.from_sequence(pendulum.period(args.start_time, args.end_time).range('hours', 12), npartitions=20)
	workers.map(calc_raob_omb).compute()
